import torch
import torchvision.transforms as transforms
import numpy as np
import cv2 as cv

from DeepSORT.ReID.nets.ResNet50 import ResNet50

# 暂存区
from config.paths import REID_MODEL

WEIGHT_PATH = REID_MODEL["WEIGHT_PATH"]


class Extractor(object):
    def __init__(self, use_cuda=True):
        self.net = ResNet50(reid=True)
        self.device = "cuda:0" if use_cuda and torch.cuda.is_available() else "cpu"
        state_dict = torch.load(WEIGHT_PATH, map_location=lambda storage, loc: storage)
        self.net.load_state_dict(state_dict)
        self.net.to(self.device)
        self.size = (64, 128)
        self.norm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def _preprocess(self, imgs_cropped):
        def _resize(image, size):
            return cv.resize(image.astype(np.float32) / 255, size)

        im_list = []
        for im in imgs_cropped:
            im = _resize(im, self.size)
            im = self.norm(im).unsqueeze(0)
            im_list.append(im)
        im_batch = torch.cat(im_list, dim=0).float()
        return im_batch

    def __call__(self, imgs_cropped):
        im_batch = self._preprocess(imgs_cropped)
        with torch.no_grad():
            im_batch = im_batch.to(self.device)
            features = self.net(im_batch)
        return features.cpu().numpy()


if __name__ == '__main__':
    from config.paths import TEST_IMAGE_PATH

    img = cv.imread(TEST_IMAGE_PATH)  # 1080, 810, 3
    # img2 = cv.imread(TEST_IMAGE_PATH)
    extr = Extractor()
    feature = extr([img])
    # feature = extr([img, img2])
    print(feature.shape)
